/*
 * @Author: 来自火星的码农 15122322+heyzhi@user.noreply.gitee.com
 * @Date: 2025-09-02 18:17:26
 * @LastEditors: 来自火星的码农 15122322+heyzhi@user.noreply.gitee.com
 * @LastEditTime: 2025-09-19 22:09:51
 * @FilePath: /server/workspace/src/base/hook.cc
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#include "hook.h"
#include "base/Log.h"
#include "base/config.h"
#include "base/fiber.h"
#include "base/iomanager.h"
#include "base/scheduler.h"
#include <asm-generic/socket.h>
#include <bits/types/struct_iovec.h>
#include <bits/types/struct_timeval.h>
#include <bits/types/timer_t.h>
#include <cerrno>
#include <cstdarg>
#include <cstddef>
#include <cstdint>
#include <ctime>
#include <fcntl.h>
#include <functional>
#include <memory>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <dlfcn.h>
#include <utility>
#include "base/timer.h"
#include "fdmanager.h"

namespace Yizhi{
    thread_local bool t_hook_enable=false;
    
    bool is_hook_enable(){
        return t_hook_enable;
    }

    void set_hook_enable(bool flag){
        t_hook_enable=flag;
    }

}
#define  HOOK_FUN(XX)\
    XX(sleep)\
    XX(usleep)\
    XX(nanosleep)\
    XX(socket)\
    XX(connect)\
    XX(accept)\
    XX(close)\
    XX(read)\
    XX(readv)\
    XX(recv)\
    XX(recvfrom)\
    XX(recvmsg)\
    XX(write)\
    XX(writev)\
    XX(send)\
    XX(sendto)\
    XX(sendmsg)\
    XX(fcntl)\
    XX(ioctl)\
    XX(setsockopt)\
    XX(getsockopt)\


    
static Yizhi::Logger::ptr g_logger=SERVER_LOG_NAME("system");
static Yizhi::ConfigVar<uint64_t>::ptr g_tcp_connect_timeout=Yizhi::Config::Lookup("tcp.connect.timeout",(uint64_t)5000,"tcp time out");
static uint64_t s_connect_timeout=-1;
void hook_init(){
    static bool is_enable=false;
    if(is_enable){
        return;
    }
#define XX(name) name ## _f=(name ## _fun)dlsym(RTLD_NEXT,#name);
    HOOK_FUN(XX);
#undef  XX
}




struct _HOOK_Int{
    _HOOK_Int(){
        hook_init();
        s_connect_timeout=g_tcp_connect_timeout->getvalue();
        g_tcp_connect_timeout->addListener([](const uint64_t&old_value,const uint64_t&new_value){
            SERVER_LOG_DEBUG(SERVER_LOG_ROOT())<<"tcp connect timeout"<<"change from old value: "<<old_value<<"to "<<new_value;
            s_connect_timeout=new_value;
        });
    }

};
static _HOOK_Int s_hook_init;

struct time_info{
    int cancel=0;
};

template<class OrginFun,typename... Args>
static ssize_t do_io(int fd,OrginFun fun,const char* hook_fun_name,
    uint64_t event,int time_so,Args&&... args){
        
        ///未启用hook
        if(!Yizhi::t_hook_enable){
            return fun(fd,std::forward<Args>(args)...);
        }

        Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);
        ///fd未创建对应ctx
        if(!ctx){
            return fun(fd,std::forward<Args>(args)...);
        }

        //已经关闭
        if(ctx->isClose()){
            errno=EBADF;
            return -1;
        }

        ///非socket文件操作符号||用户自定义为nonblock
        if(!ctx->isSocket()||ctx->getUserNonblock()){
            return fun(fd,std::forward<Args>(args)...);
        }
        
        uint64_t to=ctx->getTime(time_so);
        std::shared_ptr<time_info> timeinfo(new time_info());  
   
retry:
        int n=fun(fd,std::forward<Args>(args)...);
        while(n==-1&&errno==EINTR){
            n=fun(fd,std::forward<Args>(args)...);
        }
        if(n==-1&&errno==EAGAIN){
            Yizhi::IOManager *iom=Yizhi::IOManager::GetThis();
            Yizhi::Timer::ptr timer;
            std::weak_ptr<time_info>winfo(timeinfo);
            if(to!=(uint64_t)-1){
                iom->addCondtionTimer(to,[iom,winfo,fd,event](){
                    auto t=winfo.lock();
                    if(!t||t->cancel){
                        return ;
                    }
                    t->cancel=ETIMEDOUT;
                    iom->cancleEvent(fd, (Yizhi::IOManager::Event)event);
                },winfo);
            }
            int rt=iom->addEvent(fd,(Yizhi::IOManager::Event)event);
            if(rt==-1){
                SERVER_LOG_DEBUG(g_logger)<<hook_fun_name<<"add event("<<fd<<","<<event<<")";
                if(timer){
                    timer->cancel();
                }
                return -1;
            }
            else{
                Yizhi::Fiber::YieldToHold();
                if(timer){
                    timer->cancel();
                }
                if(timeinfo->cancel){
                    errno=timeinfo->cancel;
                    return -1;
                }
                goto retry;
            }
        }
        return n; 

    

        
}





extern "C"{
#define XX(name) name ## _fun name ## _f=nullptr;
    HOOK_FUN(XX)
#undef XX


    unsigned int sleep(unsigned int second){
        if(!Yizhi::is_hook_enable()){
            return sleep_f(second);
        }
        Yizhi::Fiber::ptr fiber=Yizhi::Fiber::GetThis();
        Yizhi::IOManager *iom=Yizhi::IOManager::GetThis();
        iom->addTimer(second*1000, [iom,fiber](){
            iom->scheduler(fiber);
        });
        Yizhi::Fiber::YieldToHold();
        return  0;

    }

    int usleep(useconds_t usec){
        if(!Yizhi::is_hook_enable()){
            return usleep_f(usec);
        }
        Yizhi::Fiber::ptr fiber=Yizhi::Fiber::GetThis();
        Yizhi::IOManager *iom=Yizhi::IOManager::GetThis();
        iom->addTimer(usec/1000, [iom,fiber](){
            iom->scheduler(fiber);
        });
        Yizhi::Fiber::YieldToHold();
        return 0;
    }
    int nanosleep(const struct timespec *requested_time, struct timespec *remaining){
        return 0;
    }

    int socket(int domain, int type, int protocol){
        if(!Yizhi::t_hook_enable){
            return socket_f(domain,type,protocol);
        }
        int fd=socket_f(domain,type,protocol);
        if(fd==-1){
            return fd;
        }
        Yizhi::FdMgr::GetInstance()->get(fd,true);
        return fd;
    }

    int connect_with_timeout(int fd,const struct sockaddr *addr,socklen_t addrlen,uint64_t timeout_ms){
        if(!Yizhi::t_hook_enable){
            return connect_f(fd,addr,addrlen);
        }
        Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);
        if(!ctx||ctx->isClose()){
            errno=EBADF;
            return -1;
        }
        if(!ctx->isSocket()){
            return connect_f(fd,addr,addrlen);
        }
        if(ctx->getUserNonblock()){
            return connect_f(fd,addr,addrlen);
        }

        int n=connect_f(fd,addr,addrlen);
        if(n==0){
            return 0;
        }

        if(n!=-1&&errno!=EINPROGRESS){
           return n;
        }


        Yizhi::IOManager* iom=Yizhi::IOManager::GetThis();
        Yizhi::Timer::ptr timer;
        std::shared_ptr<time_info> tinfo(new time_info);
        std::weak_ptr<time_info> weak_info(tinfo);

        if(timeout_ms!=(uint64_t)-1){
            timer=iom->addCondtionTimer(timeout_ms,[weak_info,iom,fd](){
                auto t=weak_info.lock();
                if(!t||t->cancel){
                    return ;
                }
                t->cancel=ETIMEDOUT;
                iom->cancleEvent(fd, Yizhi::IOManager::Event::WRITE);
            },weak_info);
        }

        int rt=iom->addEvent(fd, Yizhi::IOManager::WRITE);
        if(rt==0){
            Yizhi::Fiber::YieldToHold();
            if(timer){
                timer->cancel();
            }
            if(tinfo->cancel){
                errno=timer->cancel();
                return -1;
            }
        }
        else {
            if(timer){
                timer->cancel();
            }
            SERVER_LOG_DEBUG(g_logger)<<"add_event error"; 
        }
        int error=0;
        socklen_t len=sizeof(int);
        if(-1==getsockopt_f(fd,SOL_SOCKET,SO_ERROR,&error,&len)){
            return -1;
        }
        if(error==0){
            return 0;
        }else{
            errno=error;
            return -1;
        }

    }
    int connect(int fd, const struct sockaddr *addr, socklen_t len){
        return  connect_with_timeout(fd,addr,len,s_connect_timeout);
    }

    int accept(int fd, struct sockaddr *__restrict addr, socklen_t *__restrict addr_len){
        int s=do_io(fd,accept_f,"accept",Yizhi::IOManager::Event::WRITE,SO_RCVTIMEO,addr,addr_len);
        if(s>=0){
            Yizhi::FdMgr::GetInstance()->get(s,true);
        }
        return s;
    }


    ///read
    ssize_t read(int fd, void *buf, size_t nbytes){
        return do_io(fd, read_f,"read_fun",Yizhi::IOManager::Event::READ,SO_RCVTIMEO,buf,nbytes);
    }

    ssize_t recv(int fd, void *buf, size_t n, int flags){
        return do_io(fd,recv_f,"recv_fun",Yizhi::IOManager::Event::READ,SO_RCVTIMEO,buf,n,flags);

    }
    ssize_t readv(int fd,const struct iovec* iov,int iovcnt){
        return  do_io(fd, readv_f,"readv_fun",Yizhi::IOManager::Event::READ,SO_RCVTIMEO,iov,iovcnt);

    }
    ssize_t recvfrom(int fd, void *__restrict buf, size_t n, int flags, struct sockaddr *__restrict addr, socklen_t *__restrict addr_len){
        return do_io(fd,recvfrom_f,"recvfrom_f",Yizhi::IOManager::Event::READ,SO_RCVTIMEO,buf,n,flags,addr,addr_len);
        
    }

    ssize_t recvmsg(int fd, struct msghdr *message, int flags){
        return do_io(fd,recvmsg_f,"recvmsg_fun",Yizhi::IOManager::Event::READ,SO_RCVTIMEO,message,flags);
    }


    ///write
    ssize_t write(int fd, const void *buf, size_t n){
        return do_io(fd,write_f,"write_fun",Yizhi::IOManager::Event::WRITE,SO_SNDTIMEO,buf,n);
    }


    ssize_t writev(int fd,const struct iovec *iov,int iovcnt){
        return do_io(fd,writev_f,"write_fun",Yizhi::IOManager::Event::WRITE,SO_SNDTIMEO,iov,iovcnt);
    
    }

    ssize_t send(int fd, const void *buf, size_t n, int flags){
        return do_io(fd,send_f,"send_fun",Yizhi::IOManager::Event::WRITE,SO_SNDTIMEO,buf,n,flags);
    }

    ssize_t sendto(int fd, const void *buf, size_t n, int flags, const struct sockaddr *addr, socklen_t addr_len){
        return do_io(fd,sendto_f,"Sendto_fun",Yizhi::IOManager::Event::READ,SO_SNDTIMEO,buf,n,flags,addr,addr_len);
    }

    ssize_t sendmsg(int fd, const struct msghdr *message, int flags){
        return do_io(fd,sendmsg_f,"send_msg_f",Yizhi::IOManager::Event::WRITE,SO_SNDTIMEO,message,flags);
    }

    

    int close(int fd){
        if(!Yizhi::t_hook_enable){
            return close_f(fd);
        }
        Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);
        if(ctx){
            Yizhi::IOManager* iom=Yizhi::IOManager::GetThis();
            if(iom){
                iom->cancleAll(fd);
            }
            Yizhi::FdMgr::GetInstance()->del(fd);
        }
        return close_f(fd);
    }



    int fcntl(int fd, int cmd, .../*arg*/){
        va_list va;
        va_start(va,cmd);


        switch (cmd) {

            case F_GETFD:
           

            case F_SETFD:
           


            case F_SETFL:{
                int arg=va_arg(va,int);
                va_end(va);

                Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);
                if(!ctx||!ctx->isSocket()||ctx->isClose()){
                    return fcntl_f(fd,cmd,arg);
                }
                ctx->setUserNonblock(arg & O_NONBLOCK);
                
                //如果是初始化时系统默认设置为阻塞
                if(ctx->getSysNonblock()){
                    arg|=O_NONBLOCK;
                }
                else{
                    arg=arg&~O_NONBLOCK;
                }
                return fcntl_f(fd,cmd,arg);
            }
             break;
           
            
            case F_GETFL:{
                va_end(va);
                int arg=fcntl_f(fd,cmd);
                Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);

                if(ctx->isClose()||!ctx->isSocket()){
                    return arg;
                }
                if(ctx->getUserNonblock()){
                    return arg|=O_NONBLOCK;
                }
                else {
                    return arg&~O_NONBLOCK;
                }
            }
            break;



            case F_DUPFD:
        
            case F_SETOWN:
            
            case F_SETSIG:

            case F_SETLEASE:

            case F_NOTIFY:
#ifdef F_SETPIPE_SZ
            case F_SETPIPE_SZ:{
                
                int arg=va_arg(va,int);
                 va_end(va);
                 return fcntl_f(fd, cmd,arg);
            }
            break;
#endif

            case F_GETPIPE_SZ:{
                va_end(va);
                return fcntl_f(fd,cmd);
            }
            break;
            
            case F_GETSIG:
            
            case F_GETLEASE:

            case F_GETOWN:
            case F_SETLK:
            case F_SETLKW:
            case F_GETLK:{
                struct flock *arg=va_arg(va,flock*);
                va_end(va);
                return fcntl_f(fd,cmd,arg);
            }break;
            case F_DUPFD_CLOEXEC:
            case F_SETOWN_EX:{
                struct f_owner_exlock* arg=va_arg(va,f_owner_exlock*);
                va_end(va);
                return  fcntl_f(fd,cmd,arg);
            }break;
            case F_GETOWN_EX:
            
    
            default:
            va_end(va);
            return fcntl_f(fd,cmd);
            break;
        }

        


    }

    ///@brief
    int ioctl(int fd, unsigned long request, ...){
        va_list va;
        va_start(va, request);
        void* arg=va_arg(va, void*);
        va_end(va);


        //如果请求是清除阻塞I/O
        if(request==FIONBIO){

           //类型转换是因为 FIONBIO 请求的接口契约规定了参数必须是一个 int*。
           bool user_nonblock=!!*(int*)arg;
           Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);
           if(!ctx||!ctx->isSocket()||ctx->isClose()){
                return ioctl(fd,request,arg);
           }

           ///@brief如果参数不满足条件则实际未达到清除阻塞I/O的条件
           ctx->setUserNonblock(user_nonblock);
        }
        return ioctl_f(fd,request,arg);
        
        
    }

    int getsockopt(int fd, int level, int optname, void *__restrict optval, socklen_t *__restrict optlen){
        return getsockopt_f(fd,level,optname,optval,optlen);
    }
    int setsockopt(int fd, int level, int optname, const void *optval, socklen_t optlen){
        if(!Yizhi::t_hook_enable){
            return  setsockopt_f(fd,level,optname,optval,optlen);
        }
        if(level==SOL_SOCKET){
            if(optname==SO_RCVTIMEO||optname==SO_SNDTIMEO){
                Yizhi::FdCtx::ptr ctx=Yizhi::FdMgr::GetInstance()->get(fd);
                if(ctx){
                    struct timeval *v=(struct timeval *)optval;
                    ctx->setTime(v->tv_sec*1000+v->tv_usec/1000, optname);
                }
            }
        }
        return setsockopt_f(fd,level,optname,optval,optlen);
    }


    




}